import torch
import cv2 as cv
import math
from C3D_model import C3D
from dataset import VideoDataset
from utils import param_disturb


use_cuda = True
print('Train c3d.')
device = torch.device('cuda:0' if use_cuda else 'cpu')
model = C3D()
model.to(device)
model.load_state_dict(torch.load('c3d.pth'))
train_dataset = VideoDataset()
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True)
optimizer = torch.optim.Adam(model.parameters(), lr=1/pow(2, 10), weight_decay=0)
loss_function = torch.nn.CrossEntropyLoss()
print('Network and dataset loading complete.')
average_loss = 0
step = 0
try:
    while True:
        for i, data in enumerate(train_dataloader):
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            predict = model(images)
            train_loss = loss_function(predict, labels)
            train_loss.backward()
            if (i+1) % 8 == 0:
                optimizer.step()
                optimizer.zero_grad()
                average_loss += train_loss
            print('\rtrain loss={:.6f}'.format(train_loss), end='')
            step += 1
            if step >= 100:
                average_loss /= step
                print(f'\nAverage loss is {average_loss}.')
                average_loss = 0
                step = 0
                # torch.save(model.state_dict(), 'c3d.pth')
                param_disturb(model, var=1e-3, device=device)
        print('\nepoch finished.')
except KeyboardInterrupt as e:
    print(e)
    print('train finished.')
    torch.save(model.state_dict(), 'c3d.pth')
